#include "IONEXWriter.h"
#include <iostream>
#include <fstream>
#include <iomanip>
#include <ctime>
#include <cmath>
#include "SphericalHarmonic.h"

IONEXWriter::IONEXWriter() {
    // 初始化
}

IONEXWriter::~IONEXWriter() {
    // 清理资源
}

void IONEXWriter::setVTECGrid(const torch::Tensor& vtecGrid) {
    vtecGrid_ = vtecGrid;
    
    // 为了向后兼容性，也存入当前设置的时间点
    TimePoint currentTime(year_, month_, day_, hour_, minute_, second_);
    timeVtecGrids_[currentTime] = vtecGrid;
}

void IONEXWriter::setVTECGridForTime(const TimePoint& time, const torch::Tensor& vtecGrid) {
    timeVtecGrids_[time] = vtecGrid;
}

void IONEXWriter::setSatelliteDCBs(const std::vector<DCBResult>& satDCBs) {
    satelliteDCBs_ = satDCBs;
    
    // 为了向后兼容性，也存入当前设置的时间点
    TimePoint currentTime(year_, month_, day_, hour_, minute_, second_);
    timeSatelliteDCBs_[currentTime] = satDCBs;
}

void IONEXWriter::setSatelliteDCBsForTime(const TimePoint& time, const std::vector<DCBResult>& satDCBs) {
    timeSatelliteDCBs_[time] = satDCBs;
}

void IONEXWriter::setReceiverDCBs(const std::vector<DCBResult>& recDCBs) {
    receiverDCBs_ = recDCBs;
    
    // 为了向后兼容性，也存入当前设置的时间点
    TimePoint currentTime(year_, month_, day_, hour_, minute_, second_);
    timeReceiverDCBs_[currentTime] = recDCBs;
}

void IONEXWriter::setReceiverDCBsForTime(const TimePoint& time, const std::vector<DCBResult>& recDCBs) {
    timeReceiverDCBs_[time] = recDCBs;
}

void IONEXWriter::setMapTime(int year, int month, int day, int hour, int minute, int second) {
    year_ = year;
    month_ = month;
    day_ = day;
    hour_ = hour;
    minute_ = minute;
    second_ = second;
}

bool IONEXWriter::writeIONEXFile(const std::string& filename) {
    std::ofstream file(filename);
    if (!file.is_open()) {
        std::cerr << "无法创建IONEX文件: " << filename << std::endl;
        return false;
    }
    
    // 写入头部
    writeHeader(file);
    
    // 确定使用按时间点存储的数据还是单一网格
    if (!timeVtecGrids_.empty()) {
        // 使用不同时间点的地图
        int mapIndex = 1;
        for (const auto& [time, vtecGrid] : timeVtecGrids_) {
            writeTECMap(file, time, mapIndex++, vtecGrid);
        }
    } else if (!vtecGrid_.defined()) {
        std::cerr << "No VTEC grid data available!" << std::endl;
        return false;
    } else {
        // 向后兼容，使用单一地图
        TimePoint currentTime(year_, month_, day_, hour_, minute_, second_);
        writeTECMap(file, currentTime, 1, vtecGrid_);
    }
    
    // 写入DCB值
    writeDCBValues(file);
    
    file.close();
    return true;
}

bool IONEXWriter::writeDCBFile(const std::string& filename) {
    std::ofstream file(filename);
    if (!file.is_open()) {
        std::cerr << "无法创建DCB文件: " << filename << std::endl;
        return false;
    }
    
    file << "# DCB File generated by GNSS VTEC Estimation System" << std::endl;
    file << "# " << std::setw(4) << year_ << "-" << std::setw(2) << std::setfill('0') << month_ << "-" 
         << std::setw(2) << day_ << " " << std::setw(2) << hour_ << ":" << std::setw(2) << minute_ 
         << ":" << std::setw(2) << second_ << std::setfill(' ') << std::endl;
    file << "#" << std::endl;
    file << "# ID              DCB(ns)      StdDev(ns)" << std::endl;
    
    // 卫星DCB
    for (const auto& dcb : satelliteDCBs_) {
        file << std::left << std::setw(16) << dcb.id << std::right 
             << std::setw(12) << std::fixed << std::setprecision(3) << dcb.dcb
             << std::setw(12) << std::fixed << std::setprecision(3) << dcb.dcbSigma << std::endl;
    }
    
    // 接收机DCB
    for (const auto& dcb : receiverDCBs_) {
        file << std::left << std::setw(16) << dcb.id << std::right 
             << std::setw(12) << std::fixed << std::setprecision(3) << dcb.dcb
             << std::setw(12) << std::fixed << std::setprecision(3) << dcb.dcbSigma << std::endl;
    }
    
    file.close();
    return true;
}

void IONEXWriter::writeHeader(std::ofstream& file) {
    // 获取当前时间
    std::time_t now = std::time(nullptr);
    std::tm* tm = std::localtime(&now);
    
    // 计算地图数量
    int numMaps = !timeVtecGrids_.empty() ? timeVtecGrids_.size() : 1;
    
    // 写入头部
    file << "     1.0            IONOSPHERE MAPS     GNSS                IONEX VERSION / TYPE\n";
    file << "INS Analysis Center                                          CREATOR\n";
    
    char dateStr[20];
    std::strftime(dateStr, sizeof(dateStr), "%Y%m%d %H%M%S", tm);
    file << dateStr << " UTC                                        " << "PGM / RUN BY / DATE\n";
    
    // 写入地图信息
    file << "  " << std::setw(6) << numMaps << "                                                      " << "# OF MAPS IN FILE\n";
    file << "  NONE                                                       " << "MAPPING FUNCTION\n";
    file << "  " << std::setw(8) << ionoHeight_ << "                                                  " << "ELEVATION CUTOFF\n";
    
    // 选择一个VTEC网格确定尺寸
    torch::Tensor referenceGrid;
    if (!timeVtecGrids_.empty()) {
        referenceGrid = timeVtecGrids_.begin()->second;
    } else {
        referenceGrid = vtecGrid_;
    }
    
    if (!referenceGrid.defined()) {
        std::cerr << "No valid VTEC grid found for header!" << std::endl;
        return;
    }
    
    // 网格定义
    int numLat = referenceGrid.size(0);
    int numLon = referenceGrid.size(1);
    double latRes = 180.0 / (numLat - 1);
    double lonRes = 360.0 / numLon;
    
    file << "     " << std::setw(5) << numLat << "     " << std::setw(5) << numLon 
         << "     " << std::setw(5) << 1 << "                              " << "HGT1 / HGT2 / DHGT\n";
    
    file << "  " << std::setw(8) << 90.0 << "  " << std::setw(8) << -90.0 
         << "  " << std::setw(8) << -latRes << "                      " << "LAT1 / LAT2 / DLAT\n";
    
    file << "  " << std::setw(8) << 0.0 << "  " << std::setw(8) << 360.0 - lonRes
         << "  " << std::setw(8) << lonRes << "                      " << "LON1 / LON2 / DLON\n";
    
    file << "    " << std::setw(6) << -1 << "                                                   " << "EXPONENT\n";
}

void IONEXWriter::writeTECMap(std::ofstream& file, const TimePoint& time, int mapIndex, const torch::Tensor& vtecGrid) {
    // 写入开始地图部分
    file << "START OF TEC MAP\n";
    file << "    " << std::setw(6) << mapIndex << "                                                    " << "MAP INDEX\n";
    
    // 写入地图时间
    file << std::setw(6) << time.year << std::setw(6) << time.month << std::setw(6) << time.day 
         << std::setw(6) << time.hour << std::setw(6) << time.minute << std::setw(6) << time.second
         << "                  " << "EPOCH OF CURRENT MAP\n";
    
    // 网格定义
    int numLat = vtecGrid.size(0);
    int numLon = vtecGrid.size(1);
    double latRes = 180.0 / (numLat - 1);
    double lonRes = 360.0 / numLon;
    
    // 写入TEC值
    for (int lat_idx = 0; lat_idx < numLat; ++lat_idx) {
        double lat = 90.0 - lat_idx * latRes;
        file << "  " << std::setw(8) << lat << "  " << std::setw(8) << 0.0 << "  " << std::setw(8) << ionoHeight_ << "\n";
        
        // 每行写5个TEC值
        for (int lon_idx = 0; lon_idx < numLon; lon_idx += 5) {
            int count = std::min(5, numLon - lon_idx);
            for (int i = 0; i < count; ++i) {
                file << std::setw(5) << static_cast<int>(vtecGrid[lat_idx][lon_idx + i].item<double>() * 10) << " ";
            }
            file << "\n";
        }
    }
    
    file << "END OF TEC MAP\n";
}

void IONEXWriter::writeDCBValues(std::ofstream& file) {
    file << "START OF RMS MAP\n";
    file << "END OF RMS MAP\n";
    
    file << "START OF DCB VALUES\n";
    
    // 使用最新时间点的DCB数据
    std::vector<DCBResult> satDCBs;
    std::vector<DCBResult> recDCBs;
    
    if (!timeSatelliteDCBs_.empty()) {
        satDCBs = timeSatelliteDCBs_.rbegin()->second;
    } else {
        satDCBs = satelliteDCBs_;
    }
    
    if (!timeReceiverDCBs_.empty()) {
        recDCBs = timeReceiverDCBs_.rbegin()->second;
    } else {
        recDCBs = receiverDCBs_;
    }
    
    // 卫星DCB
    for (const auto& dcb : satDCBs) {
        file << "DSB " << dcb.id << " P1P2 " << std::fixed << std::setprecision(3) << dcb.dcb << " " 
             << std::fixed << std::setprecision(3) << dcb.dcbSigma << "\n";
    }
    
    // 接收机DCB
    for (const auto& dcb : recDCBs) {
        file << "DSB " << dcb.id << " P1P2 " << std::fixed << std::setprecision(3) << dcb.dcb << " " 
             << std::fixed << std::setprecision(3) << dcb.dcbSigma << "\n";
    }
    
    file << "END OF DCB VALUES\n";
    
    // 文件结束
    file << "END OF HEADER\n";
}

bool IONEXWriter::createIONEXFile(const std::string& filename) {
    // 保存文件名
    outputFilename_ = filename;
    
    // 创建文件流
    outputFile_ = new std::ofstream(filename);
    if (!outputFile_ || !outputFile_->is_open()) {
        std::cerr << "无法创建IONEX文件: " << filename << std::endl;
        return false;
    }
    
    // 写入头部
    writeHeader(*outputFile_);
    
    return true;
}

bool IONEXWriter::addEpoch(const TimePoint& time, 
                          const std::vector<double>& shCoeffs,
                          const std::map<std::string, double>& satDCBs,
                          const std::map<std::string, double>& recDCBs) {
    if (!outputFile_ || !outputFile_->is_open()) {
        std::cerr << "IONEX文件未打开，无法添加时间点数据" << std::endl;
        return false;
    }
    
    // 保存球谐系数
    timeShCoeffs_[time] = shCoeffs;
    
    // 保存DCB值
    timeSatDCBMap_[time] = satDCBs;
    timeRecDCBMap_[time] = recDCBs;
    
    // 计算球谐阶数
    int degree = static_cast<int>(std::sqrt(shCoeffs.size()) - 1);
    
    // 生成VTEC网格
    torch::Tensor vtecGrid = generateVTECGridFromCoefficients(shCoeffs, degree);
    
    // 保存VTEC网格
    timeVtecGrids_[time] = vtecGrid;
    
    // 转换DCB结果为向量格式
    std::vector<DCBResult> satDCBResults;
    std::vector<DCBResult> recDCBResults;
    
    for (const auto& [id, value] : satDCBs) {
        DCBResult result;
        result.id = id;
        result.dcb = value;
        result.dcbSigma = 0.1;  // 默认标准差
        satDCBResults.push_back(result);
    }
    
    for (const auto& [id, value] : recDCBs) {
        DCBResult result;
        result.id = id;
        result.dcb = value;
        result.dcbSigma = 0.1;  // 默认标准差
        recDCBResults.push_back(result);
    }
    
    // 保存DCB结果
    timeSatelliteDCBs_[time] = satDCBResults;
    timeReceiverDCBs_[time] = recDCBResults;
    
    // 写入TEC地图
    static int mapIndex = 1;
    writeTECMap(*outputFile_, time, mapIndex++, vtecGrid);
    
    return true;
}

bool IONEXWriter::finalizeIONEXFile() {
    if (!outputFile_ || !outputFile_->is_open()) {
        std::cerr << "IONEX文件未打开，无法完成写入" << std::endl;
        return false;
    }
    
    // 写入DCB值
    writeDCBValues(*outputFile_);
    
    // 关闭文件
    outputFile_->close();
    delete outputFile_;
    outputFile_ = nullptr;
    
    std::cout << "IONEX文件写入完成: " << outputFilename_ << std::endl;
    return true;
}

torch::Tensor IONEXWriter::generateVTECGridFromCoefficients(
    const std::vector<double>& shCoeffs, 
    int degree, 
    double latRes, 
    double lonRes) {
    
    // 计算网格大小
    int numLat = static_cast<int>(180.0 / latRes) + 1;
    int numLon = static_cast<int>(360.0 / lonRes);
    
    // 创建网格
    torch::Tensor vtecGrid = torch::zeros({numLat, numLon});
    
    // 创建球谐模型
    SphericalHarmonic shModel(degree);
    shModel.setCoefficients(shCoeffs);
    
    // 填充网格
    for (int i = 0; i < numLat; ++i) {
        double lat = 90.0 - i * latRes;
        for (int j = 0; j < numLon; ++j) {
            double lon = j * lonRes;
            double vtec = shModel.evaluateVTEC(lat, lon);
            vtecGrid[i][j] = vtec;
        }
    }
    
    return vtecGrid;
} 